{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Module: Data.Vector.Sized.Sample
-- Copyright: (c) 2018 Koz Ross
-- Maintainer: Koz Ross <koz.ross@retro-freedom.nz>
-- Stability: Experimental
-- Portability: Non-portable
--
-- Implementations of various sampling algorithms, both weighted and unweighted.
-- These return new copies, and guarantee ordering stability (that is,
-- the sampled elements will be in the same order that they were in the array
-- they were taken from).
module Data.Vector.Sized.Sample where

import Control.Monad.Random.Strict
import Control.Monad.ST
import Data.Foldable
import Data.Proxy
import GHC.TypeNats

import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable.Sized as VGMS
import qualified Data.Vector.Generic.Sized as VGS

-- | Random sampling using Algorithm R.
--
-- __Complexity:__ \(O(n)\)
{-# INLINE sample #-}
sample ::
     forall k m v n a l.
     (KnownNat n, KnownNat k, MonadRandom m, VG.Vector v a, (k + l) ~ n)
  => VGS.Vector v n a
  -> m (VGS.Vector v k a)
sample v = do
  ixes <- traverse (\i -> (i, ) <$> getRandomR (0, i)) [k' .. (n' - 1)]
  let v' =
        runST
          (do vec <- VGMS.unsafeNew
              VGS.copy vec (VGS.take @_ @k v)
              traverse_
                (\(i, j) ->
                   when (j < k') (VGMS.unsafeWrite vec j (VGS.unsafeIndex v i)))
                ixes
              VGS.unsafeFreeze vec)
  return v'
  where
    k' = fromIntegral . natVal $ Proxy @k
    n' = VGS.length v

-- | Weighted random sampling using Chao's algorithm.
--
-- The algorithm is based on the following approach to weighting: let the weight
-- of two items \(i,j\) be \(w_i,w_j\), and \(p_i\) be the probability that
-- \(i\) is selected; then \(p_j = \min\{1, p_i \cdot \frac{w_j}{w_i}\}\).
--
-- __Complexity:__ \(O(n)\)
{-# INLINE chao #-}
chao ::
     forall k m v n a b l.
     ( KnownNat n
     , KnownNat k
     , MonadRandom m
     , Fractional b
     , Ord b
     , Random b
     , VG.Vector v (a, b)
     , VG.Vector v a
     , (k + l) ~ n
     )
  => VGS.Vector v n (a, b)
  -> m (VGS.Vector v k a)
chao v = do
  ixes <-
    traverse
      (\i -> (i, , ) <$> getRandomR (0.0, 1.0) <*> getRandomR (0, k' - 1))
      [k' .. (n' - 1)]
  let v' =
        runST
          (do vec <- VGMS.unsafeNew
              wSum <-
                foldlM
                  (\s i -> do
                     let (e, w) = VGS.unsafeIndex v i
                     VGMS.unsafeWrite vec i e
                     return (s + w / fromIntegral k'))
                  0.0
                  [0 .. (k' - 1)]
              foldM_
                (\s (i, j, k'') -> do
                   let (e, w) = VGS.unsafeIndex v i
                   let s' = s + w / fromIntegral k'
                   let p = w / s'
                   when (j <= p) (VGMS.unsafeWrite vec k'' e)
                   return s')
                wSum
                ixes
              VGS.unsafeFreeze vec)
  return v'
  where
    k' = fromIntegral . natVal $ Proxy @k
    n' = VGS.length v
